# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：SC_half_plane_change_floor_corner.py
    Created on：2025/6/4 13:53
    Author：(input)
    Description: 倾斜壁面半无限大房间势流源模型，地面拐角变化floor
"""
import numpy as np
import matplotlib.pyplot as plt


class InclinedWallPotentialPlotter:
    def __init__(self, H, L, resolution, jiaodu=0.0, Q_s=1.0, S=0.2, bottom_x_start=0.0):
        self.h = H
        self.L = L
        alpha = jiaodu * (np.pi / 180)  # 角度转换为弧度
        self.alpha = alpha
        self.resolution = resolution
        self.Q_s = Q_s
        self.bottom_x_start = bottom_x_start

        if jiaodu < 0:
            self.x_c = S - H * np.tan(alpha)
        else:
            self.x_c = S + H * np.tan(alpha)  # 送风中心的水平坐标
        self.x_min = min(0.0, H * np.tan(alpha))  # 左边界最小 x 坐标

    def generate_grid(self):
        """
        生成复平面网格 z = x + i y，确保所有点位于倾斜壁面右侧（x >= y * tan(alpha) + bottom_x_start）。
        返回未倾斜坐标 (X, Y)、倾斜坐标 (X_tilted, Y)、复数网格 Z。
        """
        # 生成 y 坐标
        y = np.linspace(0, self.h, self.resolution)

        # 计算每层 y 对应的倾斜壁面 x 坐标
        x_wall = y * np.tan(self.alpha) + self.bottom_x_start

        # 初始化网格数组
        X = np.zeros((self.resolution, self.resolution))
        Y = np.zeros((self.resolution, self.resolution))

        # 为每个 y_i 生成 x 坐标，从 x_wall[i] 到 L
        for i in range(self.resolution):
            x_min_i = x_wall[i]  # 倾斜壁面处 x 坐标
            x = np.linspace(x_min_i, self.L, self.resolution)
            X[i, :] = x
            Y[i, :] = y[i]

        # 计算倾斜坐标，调整底点偏移
        X_tilted = X - self.bottom_x_start
        Z = X_tilted + 1j * Y

        return X, Y, X_tilted, Z

    def compute_complex_potential(self, Z):
        """
        计算复势函数 w(z) = φ + iψ，并返回 φ 和 ψ。
        """
        kv = self.Q_s / np.pi
        A = np.pi / self.h
        B = 1 - np.tan(self.alpha)
        C = 1 - np.tan(self.alpha) ** 2

        X = Z.real
        Y = Z.imag

        Rex = (X * B + Y * np.tan(self.alpha)) / C
        Imx = Y / -C
        F1 = np.cosh(A * Rex) * np.cos(A * Imx) + 1j * (np.sinh(A * Rex) * np.sin(A * Imx))

        Imc = self.h / -C
        Rec = (self.x_c * B + self.h * np.tan(self.alpha)) / C
        F2 = np.cosh(A * Rec) * np.cos(A * Imc) + 1j * (np.sinh(A * Rec) * np.sin(A * Imc))

        F = (np.cosh(A * Rex) * np.cos(A * Imx) - np.cosh(A * Rec) * np.cos(A * Imc)) \
            + 1j * (np.sinh(A * Rex) * np.sin(A * Imx) - np.sinh(A * Rec) * np.sin(A * Imc))

        W = kv * np.log(F)

        phi = np.real(W)
        psi = np.imag(W)

        return phi, psi

    def compute_velocity_field(self, psi, X_tilted, Y):
        """
        计算速度场 (u, v) 基于流函数 ψ，应用边界遮罩。
        u = ∂ψ/∂y, v = -∂ψ/∂x
        """
        dy = Y[1, 0] - Y[0, 0]
        u = np.gradient(psi, dy, axis=0)

        dx = X_tilted[0, 1] - X_tilted[0, 0]
        v = -np.gradient(psi, dx, axis=1)

        # 遮罩：仅保留物理边界内的速度
        mask = X_tilted >= Y * np.tan(self.alpha) + self.bottom_x_start
        u[~mask] = 0
        v[~mask] = 0

        return u, v

    def plot_physical_wall_boundaries(self, ax):
        """
        绘制倾斜壁面与边界。
        """
        y_vals = np.linspace(0, self.h, 200)
        x_left = y_vals * np.tan(self.alpha) + self.bottom_x_start
        ax.plot(x_left, y_vals, 'k-', linewidth=2.5, label='Inclined Wall')

        top_x_start = y_vals[-1] * np.tan(self.alpha) + self.bottom_x_start
        x_top_right = np.linspace(top_x_start, self.L, 200)
        ax.plot(x_top_right, [self.h] * len(x_top_right), 'k-', linewidth=2.5)

        ax.plot([self.bottom_x_start, self.L], [0, 0], 'k-', linewidth=2.5)

        if abs(self.alpha) > 1e-8:
            ax.plot([self.bottom_x_start, self.bottom_x_start], [0, self.h], 'k--', linewidth=1.5)
            ax.plot([self.bottom_x_start, top_x_start], [self.h, self.h], 'k--', linewidth=1.5)

        ax.scatter(self.x_c, self.h, color='r', s=80, marker='o', label='Source (Inlet)')

    def plot_valid_grid_points(self):
        """
        可视化复平面 Z 中的有效网格点在物理坐标系下的位置。
        """
        X, Y, X_tilted, Z = self.generate_grid()

        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(X_tilted.flatten(), Y.flatten(), 'b.', markersize=0.5)

        self.plot_physical_wall_boundaries(ax)

        x_margin_ratio = 0.05
        y_margin_ratio = 0.02
        x_range = self.L - self.x_min
        y_range = self.h
        x_margin = x_range * x_margin_ratio
        y_margin = y_range * y_margin_ratio

        ax.set_xlim(self.x_min - x_margin, self.L + x_margin)
        ax.set_ylim(0 - y_margin, self.h + y_margin)

        ax.set_title("Valid Grid Physical Plane")
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect('equal')
        ax.grid(True)
        plt.tight_layout()
        plt.show()

    def plot_velocity_field(self, show_magnitude=False, title='Velocity Field with Inclined Wall', cmap='jet'):
        """
        绘制速度场流线图，用箭头表示方向，可选通过颜色和颜色条表示速度大小。
        """
        X, Y, X_tilted, Z = self.generate_grid()
        phi, psi = self.compute_complex_potential(Z)
        u, v = self.compute_velocity_field(psi, X_tilted, Y)

        if not (np.issubdtype(u.dtype, np.number) and np.issubdtype(v.dtype, np.number)):
            raise ValueError("u and v must contain numeric values")
        if np.min(u) == np.max(u) and np.min(v) == np.max(v):
            raise ValueError("u and v must have a range of values for stream plotting")

        magV = None
        if show_magnitude:
            magV = np.sqrt(u**2 + v**2)
            if not np.issubdtype(magV.dtype, np.number):
                raise ValueError("magV must contain numeric values")
            if np.min(magV) == np.max(magV):
                raise ValueError("magV must have a range of values for color mapping")

        fig, ax = plt.subplots(figsize=(10, 5))
        if show_magnitude:
            stream = ax.streamplot(X, Y, u, v,
                                   density=1, linewidth=1,
                                   arrowsize=1, arrowstyle='->',
                                   color=magV, cmap=cmap)
            plt.colorbar(stream.lines, label='Velocity Magnitude (m/s)')
        else:
            stream = ax.streamplot(X, Y, u, v,
                                   density=1, linewidth=1,
                                   arrowsize=1, arrowstyle='->')

        self.plot_physical_wall_boundaries(ax)

        x_margin_ratio = 0.05
        y_margin_ratio = 0.02
        x_range = self.L - self.x_min
        y_range = self.h
        x_margin = x_range * x_margin_ratio
        y_margin = y_range * y_margin_ratio

        ax.set_xlim(self.x_min - x_margin, self.L + x_margin)
        ax.set_ylim(0 - y_margin, self.h + y_margin)

        ax.set_xlabel("width")
        ax.set_ylabel("height")
        ax.set_title(title)
        ax.set_aspect('equal')
        ax.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()

        x_flat = X_tilted.flatten()
        y_flat = Y.flatten()
        u_flat = u.flatten()
        v_flat = v.flatten()

        data = np.column_stack((x_flat, y_flat, u_flat, v_flat))
        np.savetxt('velocity_data.csv', data, delimiter=',',
                   header='x,y,u,v', comments='')

    def plot_physical_quantities_along_line(self, x_fixed, y_range=None):
        """
        指定固定 x 坐标和 y 范围，计算并绘制 phi, psi, u, v, magV 的线图。
        横轴为 y，纵轴为各物理量。

        Parameters:
        ----------
        x_fixed : float
            固定的 x 坐标值
        y_range : tuple, optional
            y 坐标范围 (ymin, ymax)，默认 [0, self.h]
        """
        if y_range is None:
            y_range = (0, self.h)
        y = np.linspace(y_range[0], y_range[1], self.resolution)
        x = np.full_like(y, x_fixed)  # 固定 x 值

        # 生成网格并计算复势
        X, Y = np.meshgrid(x, y)
        X_tilted = X + Y * np.tan(self.alpha) - self.bottom_x_start
        Z = X_tilted + 1j * Y

        phi, psi = self.compute_complex_potential(Z)

        # 计算速度场
        dy = y[1] - y[0]
        u = np.gradient(psi, dy, axis=0)
        v = np.zeros_like(u)  # v 偏导数在此1D切片中为0，因为x固定
        magV = np.sqrt(u**2 + v**2)

        # 遮罩：仅保留物理边界内的值
        mask = X_tilted >= y * np.tan(self.alpha) + self.bottom_x_start
        phi[~mask] = np.nan
        psi[~mask] = np.nan
        u[~mask] = np.nan
        v[~mask] = np.nan
        magV[~mask] = np.nan

        # 绘制线图
        fig, axs = plt.subplots(5, 1, figsize=(10, 15), sharex=True)
        titles = ['Potential φ', 'Stream Function ψ', 'Velocity u', 'Velocity v', 'Velocity Magnitude']
        data = [phi, psi, u, v, magV]

        for ax, title, d in zip(axs, titles, data):
            ax.plot(y, d, 'b-')
            ax.set_title(title)
            ax.grid(True)
            ax.set_ylabel(f'{title} (unitless)')  # 单位可根据需要调整

        axs[-1].set_xlabel('y (m)')
        plt.tight_layout()
        plt.show()

    def plot(self):
        """
        绘制 ψ(x, y) 和壁面结构，包含适当边距，并保存数据到 CSV。
        """
        X, Y, X_tilted, Z = self.generate_grid()
        phi, psi = self.compute_complex_potential(Z)

        fig, ax = plt.subplots(figsize=(10, 5))

        psi_levels = [0, 1]
        cs = ax.contour(X_tilted, Y, psi, levels=20, colors='blue', linewidths=1)
        ax.clabel(cs, inline=True, fontsize=7, fmt='ψ')

        self.plot_physical_wall_boundaries(ax)

        x_margin_ratio = 0.05
        y_margin_ratio = 0.02
        x_range = self.L - self.x_min
        y_range = self.h
        x_margin = x_range * x_margin_ratio
        y_margin = y_range * y_margin_ratio

        ax.set_xlim(self.x_min - x_margin, self.L + x_margin)
        ax.set_ylim(0 - y_margin, self.h + y_margin)

        ax.set_title("Stream Function ψ(x, y) with Inclined Wall")
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect('equal')
        ax.grid(True)
        plt.tight_layout()
        plt.show()

        x_flat = X_tilted.flatten()
        y_flat = Y.flatten()
        phi_flat = phi.flatten()
        psi_flat = psi.flatten()

        data = np.column_stack((x_flat, y_flat, phi_flat, psi_flat))
        np.savetxt('potential_data.csv', data, delimiter=',',
                   header='x,y,phi,psi', comments='')

        self.plot_velocity_field(show_magnitude=True)




# 示例用法
if __name__ == "__main__":
    resolution = 10
    U_0 = -20
    b = 0.02
    jiaodu = -2
    H = 5.0
    L = 10.0
    S = 1.3
    Q_s = U_0 * b
    bottom_x_start = 0.0

    plotter = InclinedWallPotentialPlotter(H, L, resolution, jiaodu, Q_s, S, bottom_x_start)
    # plotter.plot_valid_grid_points()
    plotter.plot()
    plotter.plot_physical_quantities_along_line(x_fixed=0.0)  # 示例：x=0, y=[0, h]